import cv2
import numpy as np

# bug制造中，暂不可用

class ChessboardDetector:
    def __init__(self, black_threshold=135, white_threshold=150):
        self.BLACK_THRESHOLD = black_threshold
        self.WHITE_THRESHOLD = white_threshold

    def preprocess(self, img):
        """图像预处理"""
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        blurred = cv2.GaussianBlur(gray, (5, 5), 1)
        _, threshold = cv2.threshold(blurred, 120, 255, cv2.THRESH_BINARY)
        return threshold

    def detect_contours(self, img):
        """检测棋盘轮廓"""
        contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if not contours:
            return None

        biggest = None
        max_area = 0

        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area > 5000:
                perimeter = cv2.arcLength(cnt, True)
                approx = cv2.approxPolyDP(cnt, 0.02 * perimeter, True)
                if len(approx) == 4 and area > max_area:
                    biggest = approx
                    max_area = area

        return biggest

    def find_key_points(self, biggest):
        """提取棋盘关键点"""
        if biggest is None or biggest.size == 0:
            return None

        points = [tuple(point[0]) for point in biggest]
        RA, RB, RC, RD = points[0], points[1], points[2], points[3]

        E = (
            (RA[0] + RB[0] + RC[0] + RD[0]) / 4,
            (RA[1] + RB[1] + RC[1] + RD[1]) / 4
        )

        key_points = {
            'RA': RA,
            'RB': RB,
            'RC': RC,
            'RD': RD,
            'E': E,
            'A': (
                RA[0] + (E[0] - RA[0]) / 3,
                RA[1] + (E[1] - RA[1]) / 3
            ),
            'C': (
                RB[0] + (E[0] - RB[0]) / 3,
                RB[1] + (E[1] - RB[1]) / 3
            ),
            'G': (
                RC[0] + (E[0] - RC[0]) / 3,
                RC[1] + (E[1] - RC[1]) / 3
            ),
            'I': (
                RD[0] + (E[0] - RD[0]) / 3,
                RD[1] + (E[1] - RD[1]) / 3
            ),
            'B': (
                (RA[0] + RB[0]) / 2,
                (RA[1] + RB[1]) / 2
            ),
            'F': (
                (RB[0] + RC[0]) / 2,
                (RB[1] + RC[1]) / 2
            ),
            'D': (
                (RC[0] + RD[0]) / 2,
                (RC[1] + RD[1]) / 2
            ),
            'H': (
                (RD[0] + RA[0]) / 2,
                (RD[1] + RA[1]) / 2
            )
        }

        return key_points

    def get_color(self, img, point, window_size=5):
        """获取指定点的颜色信息"""
        half_window = window_size // 2
        px, py = point

        # 将坐标转换为整数
        px = int(px)
        py = int(py)

        # 检查是否超出图像范围
        img_height, img_width, _ = img.shape
        if px < 0 or px >= img_width or py < 0 or py >= img_height:
            return 0

        black_count = 0
        white_count = 0

        for i in range(-half_window, half_window + 1):
            for j in range(-half_window, half_window + 1):
                x = px + i
                y = py + j

                if 0 <= x < img_width and 0 <= y < img_height:
                    b, g, r = img[y, x]
                    if r < self.BLACK_THRESHOLD and g < self.BLACK_THRESHOLD and b < self.BLACK_THRESHOLD:
                        black_count += 1
                    elif r > self.WHITE_THRESHOLD and g > self.WHITE_THRESHOLD and b > self.WHITE_THRESHOLD:
                        white_count += 1

        if black_count > white_count:
            return 1
        elif white_count > black_count:
            return 2
        else:
            return 0

    def detect(self, img):
        """检测棋盘并返回结果"""
        processed = self.preprocess(img)
        biggest = self.detect_contours(processed)
        if biggest is None:
            return {'message': 'No chessboard detected'}

        key_points = self.find_key_points(biggest)
        if key_points is None:
            return {'message': 'No key points found'}

        colors = {}
        for key, point in key_points.items():
            color = self.get_color(img, point)
            colors[key] = color

        return colors

# 示例代码
if __name__ == "__main__":
    cap = cv2.VideoCapture(0)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

    detector = ChessboardDetector()

    while True:
        ret, frame = cap.read()
        if not ret: break

        results = detector.detect(frame)

        # 输出检测结果
        print("Detected colors:", results)

        cv2.imshow("Frame", frame)
        if cv2.waitKey(1) == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()